# -*- coding:utf-8 -*-
"""
author: liuxu
date: 2025.2.12
description: 批量编译cpp文件，为后续hipprof测试使用
"""

import os
import subprocess
import traceback
import concurrent.futures
from common.logger_ctrl import mylogger
from setting import CPP_DIR, EXE_DIR


# cpp用例：（源码路径，编译后可执行文件路径）
cpp_case_map = {
    "matrix": "matrix.cpp",  # 矩阵转置测例调用7个hip-api（包括隐式调用），1个kernel
    "multi_devices_test": "multi_devices_test.cpp",
    "roctracer_test": "roctracer_test.cpp",
    "roctx_test": "roctx_test.cpp",
    "multiprocess_test": "multiprocess_test.cpp",
    "multiprocess_multidevice_test": "multiprocess_multidevice_test.cpp",
    "cpu_memory_leak_test": "cpu_memory_leak_test.cpp",
    "gpu_memory_leak_test": "gpu_memory_leak_test.cpp",
    "stream_and_event_leak_test": "stream_event_leak_test.cpp",
    "virtual_memory_leak_test": "virtual_memory_leak_test.cpp",
}

def compile_cpp_file(src_file_path, target_file_path, *compile_args, compiler="dcc"):
    """
    使用hipcc编译c++文件至目标路径
    :param src_file_path: 源码全路径
    :param target_file_path: 编译后可执行文件路径
    :param compile_args: 编译参数
    :param compiler: 编译工具hipcc/dcc/nvcc/flang
    :return:
    """
    # 检查可执行文件是否已存在，存在则跳过编译
    if os.path.exists(target_file_path):
        mylogger.info(f"'{target_file_path}' already exists,skip compile.")
        return True
    # 检查输出目录
    out_dir = os.path.dirname(target_file_path)
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)
    compile_command = [compiler, "-o", target_file_path, src_file_path, *compile_args]
    mylogger.info(f"Start to exec compile command: {' '.join(compile_command)}")
    try:
        ret = subprocess.run(compile_command, check=True)  # 返回非零退出状态码时会引发 CalledProcessError 异常
        mylogger.info(f"Succeed to compile {src_file_path}!")
        return ret.returncode == 0
    except subprocess.CalledProcessError as e:
        mylogger.error(f"Failed to compile {src_file_path}, an error occurred during compilation: {e}\n{traceback.format_exc()}")
    return False

def compile_cpp_files_concurrently():
    """
    多线程编译多个 C++ 文件
    :return: 编译成功和失败的文件列表
    """

    def gen_cpp_file_map(src_dir, exe_dir):
        cpp_file_map = {} ## 读取cpp文件，转为为{/src_path/to/demo.cpp:/build_path/to/demo,}格式

        # 遍历cpp所在目录，查找所有的.cpp文件
        for file in os.listdir(src_dir):
            if file.endswith(".cpp"):
                # 生成源文件的完整路径
                src_file_path = os.path.join(src_dir, file)
                # 生成对应的可执行文件路径
                exe_file_name = os.path.splitext(file)[0]  # 去掉扩展名
                exe_file_path = os.path.join(exe_dir, exe_file_name)
                cpp_file_map[src_file_path] = exe_file_path
        return cpp_file_map
    # 存储编译成功和失败的文件
    successful_compiles = []
    failed_compiles = []
    cpp_to_exe_file_map = gen_cpp_file_map(CPP_DIR, EXE_DIR)
    # print(cpp_to_exe_file_map)
    # 创建线程池
    with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()*2) as executor:
        # 提交编译任务到线程池
        future_to_file = {executor.submit(compile_cpp_file, cpp_file, exe_file): cpp_file for cpp_file, exe_file in cpp_to_exe_file_map.items()}
        # 处理每个任务的结果
        for future in concurrent.futures.as_completed(future_to_file):
            cpp_file = future_to_file[future]
            try:
                if future.result() is True:  # 判断线程返回结果
                    successful_compiles.append(cpp_file)
                else:
                    failed_compiles.append(cpp_file)
            except Exception as e:
                mylogger.error(f'There was an exception in compilation {cpp_file}:\n {e}')
    return successful_compiles, failed_compiles


if __name__ == '__main__':
    compile_cpp_files_concurrently()